from pysimm import system, lmps, forcefield
from pysimm.apps.random_walk import random_walk
from pva import monomer
from pvahead import header
from pvatail import tailer

def run(test=False):
    # we'll create a pe monomer from the pysimm.models database
    pe = monomer()
    # add head and tail
    head = header()
    tail = tailer()  
    # we'll instantiate a GAFF2 forcefield object for use later
    f = forcefield.Gaff2()
    
    # the monomers do not have any charges, so we will derive partial charges using the gasteiger algorithm
    pe.apply_charges(f, charges='gasteiger')
    head.apply_charges(f, charges='gasteiger')
    tail.apply_charges(f, charges='gasteiger')
    
    # then we use the pattern argument to define how many of each "monomers" to add. Let's add 5 more monomers to our chain
    
    # first import the copolymer function
    
    from pysimm.apps.random_walk import copolymer
    
    longer_polymer = copolymer([head, pe, tail], 12, pattern=[1, 10, 1], forcefield=f)
    longer_polymer.write_xyz('longer_polymer.xyz')
    longer_polymer.write_yaml('longer_polymer.yaml')
    longer_polymer.write_lammps('longer_polymer.lmps')
    longer_polymer.write_chemdoodle_json('longer_polymer.json')

if __name__ == '__main__':
    run()